/*
* Copyright (c) Huawei Technologies Co., Ltd. 2020-2021. All rights reserved.
* gazelle is licensed under the Mulan PSL v2.
* You can use this software according to the terms and conditions of the Mulan PSL v2.
* You may obtain a copy of Mulan PSL v2 at:
*     http://license.coscl.org.cn/MulanPSL2
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
* PURPOSE.
* See the Mulan PSL v2 for more details.
*/

#include <stdlib.h>
#include <string.h>
#include <securec.h>
#include <arpa/inet.h>

#include <lwip/lwipgz_hlist.h>

#include <rte_rwlock.h>
#include <rte_malloc.h>
#include <rte_memzone.h>
#include <rte_jhash.h>
#include <rte_lcore.h>

#include "lstack_log.h"
#include "common/gazelle_base_func.h"
#include "lstack_sock_htable.h"

#define SOCK_HTABLE_HASH_NAME       "lstack_sock_htable"

static struct sock_htable *g_sock_htable = NULL;

struct sock_htable *sock_htable_create(uint32_t max_sock_num)
{
    struct sock_htable *sock_htable = NULL;
    char name[RTE_MEMZONE_NAMESIZE] = {0};
    const struct rte_memzone *mz;

    snprintf_s(name, sizeof(name), sizeof(name) - 1, SOCK_HTABLE_HASH_NAME);
    mz = rte_memzone_reserve(name, sizeof(struct sock_htable), rte_socket_id(), 0);
    if (mz == NULL) {
        LSTACK_LOG(ERR, LSTACK, "rte_memzone_reserve failed\n");
        return NULL;
    }
    memset_s(mz->addr, mz->len, 0, mz->len);
    sock_htable = mz->addr;

    rte_rwlock_init(&sock_htable->rwlock);

    for (int i = 0; i < GAZELLE_MAX_TCP_SOCK_HTABLE_SIZE; i++) {
        hlist_init_head(&sock_htable->array[i].chain);
        sock_htable->array[i].chain_size = 0;
    }
    sock_htable->cur_sock_num = 0;
    sock_htable->max_sock_num = max_sock_num;

    g_sock_htable = sock_htable;

    return sock_htable;
}

struct sock_htable *sock_htable_attach(void)
{
    const struct rte_memzone *mz = rte_memzone_lookup(SOCK_HTABLE_HASH_NAME);
    if (mz == NULL) {
        LSTACK_LOG(ERR, LSTACK, "sock htable attach failed\n");
        return NULL;
    }
    g_sock_htable = mz->addr;

    printf("sock_htable_attach success!!!\n");

    return g_sock_htable;
}

void sock_htable_destroy(void)
{
    struct sock_htable_node *sock_node = NULL;
    struct hlist_node *node = NULL;
    struct sock_htable *sock_htable = g_sock_htable;
    const struct rte_memzone *mz;
    uint32_t i;

    if (sock_htable == NULL) {
        return;
    }

    for (i = 0; i < GAZELLE_MAX_TCP_SOCK_HTABLE_SIZE; i++) {
        node = sock_htable->array[i].chain.first;
        while (node != NULL) {
            sock_node = hlist_entry(node, typeof(*sock_node), hlist_node);
            node = node->next;
            hlist_del_node(&sock_node->hlist_node);
            rte_free(sock_node);
        }
    }

    mz = rte_memzone_lookup(SOCK_HTABLE_HASH_NAME);
    if (mz != NULL) {
        rte_memzone_free(mz);
    }
}

static __rte_always_inline uint32_t sock_htable_hash_fn(void *key, int len)
{
    return rte_jhash(key, len, 0);
}

static struct sock_htable_bucket *sock_htable_bucket_get(struct sock_htable *sock_htable, struct sock_htable_key *key)
{
    uint32_t index = sock_htable_hash_fn(key, sizeof(struct sock_htable_key)) % GAZELLE_MAX_TCP_SOCK_HTABLE_SIZE;
    return &sock_htable->array[index];
}

int sock_htable_add_key_data(struct sock_htable_key *key, uint16_t queue_id)
{
    struct sock_htable *sock_htable = g_sock_htable;
    struct sock_htable_node *sock_node= NULL;
    struct hlist_node *node = NULL;
    struct hlist_head *head = NULL;
    struct sock_htable_bucket *sock_bucket = NULL;

    sock_bucket = sock_htable_bucket_get(sock_htable, key);
    if (sock_bucket == NULL) {
        LSTACK_LOG(ERR, LSTACK, "sock_htable_bucket_get failed\n");
        return -1;
    }

    rte_rwlock_read_lock(&sock_htable->rwlock);
    head = &sock_bucket->chain;
    hlist_for_each_entry(sock_node, node, head, hlist_node) {
        if ((memcmp(&sock_node->key, key, sizeof(struct sock_htable_key)) == 0) && sock_node->queue_id == queue_id) {
            rte_rwlock_read_unlock(&sock_htable->rwlock);
            LSTACK_LOG(ERR, LSTACK, "sock htable node exist, queue: %d\n", queue_id);
            return -1;
        }
    }
    rte_rwlock_read_unlock(&sock_htable->rwlock);

    if (sock_htable->cur_sock_num == sock_htable->max_sock_num) {
        LSTACK_LOG(ERR, LSTACK, "sock htable has readched the maximum\n");
        return -1;
    }

    // TODO: not use rte_malloc
    sock_node = rte_malloc(NULL, sizeof(struct sock_htable_node), RTE_CACHE_LINE_SIZE);
    if (sock_node == NULL) {
        LSTACK_LOG(ERR, LSTACK, "rte_malloc failed\n");
        return -1;
    }
    memcpy(&sock_node->key, key, sizeof(struct sock_htable_node));
    sock_node->queue_id = queue_id;
    sock_node->cur_conn = 0;

    rte_rwlock_write_lock(&sock_htable->rwlock);
    hlist_add_head(&sock_node->hlist_node, &sock_bucket->chain);
    sock_htable->cur_sock_num++;
    sock_bucket->chain_size++;
    rte_rwlock_write_unlock(&sock_htable->rwlock);

    return 0;
}

void sock_htable_del_by_queue(uint16_t queue_id)
{
    struct sock_htable *sock_htable = g_sock_htable;
    struct sock_htable_node *sock_node = NULL;
    struct hlist_node *node = NULL;
    struct hlist_head *head = NULL;
    struct sock_htable_bucket *sock_bucket = NULL;

    rte_rwlock_write_lock(&sock_htable->rwlock);
    for (int i = 0; i < GAZELLE_MAX_TCP_SOCK_HTABLE_SIZE; i++) {
        sock_bucket = &sock_htable->array[i];
        if (sock_bucket->chain_size == 0) {
            continue;
        }
        head = &sock_bucket->chain;
        hlist_for_each_entry(sock_node, node, head, hlist_node) {
            if (sock_node->queue_id == queue_id) {
                hlist_del_node(&sock_node->hlist_node);
                sock_htable->cur_sock_num--;
                sock_bucket->chain_size--;
                rte_free(sock_node);
            }
        }
    }
    rte_rwlock_write_unlock(&sock_htable->rwlock);
}

int sock_htable_del_key(struct sock_htable_key *key, uint16_t queue_id)
{
    struct sock_htable *sock_htable = g_sock_htable;
    struct sock_htable_node *sock_node = NULL;
    struct hlist_node *node = NULL;
    struct hlist_head *head = NULL;
    struct sock_htable_bucket *sock_bucket = NULL;

    sock_bucket = sock_htable_bucket_get(sock_htable, key);
    if (sock_bucket == NULL) {
        LSTACK_LOG(ERR, LSTACK, "sock_htable_bucket_get failed\n");
        return -1;
    }

    rte_rwlock_write_lock(&sock_htable->rwlock);
    head = &sock_bucket->chain;
    hlist_for_each_entry(sock_node, node, head, hlist_node) {
        if ((memcmp(&sock_node->key, key, sizeof(struct sock_htable_key)) == 0) && (sock_node->queue_id == queue_id)) {
            break;
        }
    }

    if (sock_node == NULL) {
        rte_rwlock_write_unlock(&sock_htable->rwlock);
        LSTACK_LOG(ERR, LSTACK, "can't find sock_node\n");
        return -1;
    }

    hlist_del_node(&sock_node->hlist_node);
    sock_htable->cur_sock_num--;
    sock_bucket->chain_size--;
    rte_rwlock_write_unlock(&sock_htable->rwlock);

    rte_free(sock_node);
    return 0;
}

int sock_htable_get_by_min_conn(struct sock_htable_key *key, uint16_t cur_queue_id, bool create)
{
    struct sock_htable *sock_htable = g_sock_htable;
    struct sock_htable_bucket *sock_bucket = NULL;
    struct sock_htable_node *sock_node_tmp = NULL;
    struct sock_htable_node *sock_node = NULL;
    struct hlist_node *node = NULL;
    struct hlist_head *head = NULL;
    uint32_t min_con = 20000;

    sock_bucket = sock_htable_bucket_get(sock_htable, key);
    if (sock_bucket == NULL) {
        LSTACK_LOG(ERR, LSTACK, "sock_htable_bucket_get failed\n");
        return -1;
    }

    rte_rwlock_read_lock(&sock_htable->rwlock);
    head = &sock_bucket->chain;
    hlist_for_each_entry(sock_node, node, head, hlist_node) {
        /* because bucket=hash(ip,port)%array_size tcp_sock maybe have other ip.port */
        if (memcmp(&sock_node->key, key, sizeof(struct sock_htable_key)) != 0) {
            continue;
        }
        if (sock_node->queue_id == cur_queue_id) {
            sock_node_tmp = sock_node;
            break;
        }
        if (sock_node->cur_conn < min_con) {
            sock_node_tmp = sock_node;
            min_con = sock_node->cur_conn;
        }
    }

    rte_rwlock_read_unlock(&sock_htable->rwlock);

    rte_rwlock_write_lock(&sock_htable->rwlock);
    if (sock_node_tmp != NULL) {
        if (create) {
            sock_node_tmp->cur_conn++;
        } else {
            sock_node_tmp->cur_conn--;
        }
        rte_rwlock_write_unlock(&sock_htable->rwlock);
        return sock_node_tmp->queue_id;
    } else {
        rte_rwlock_write_unlock(&sock_htable->rwlock);
        return -1;
    }
}

void sock_htable_dump(void)
{
    struct sock_htable *sock_htable = g_sock_htable;
    struct sock_htable_node *sock_node = NULL;
    struct sock_htable_bucket *sock_bucket = NULL;
    struct hlist_node *node = NULL;
    struct hlist_head *head = NULL;
    char addr[16];

    rte_rwlock_read_lock(&sock_htable->rwlock);
    for (int i = 0; i < GAZELLE_MAX_TCP_SOCK_HTABLE_SIZE; i++) {
        sock_bucket = &sock_htable->array[i];
        head = &sock_bucket->chain;
        hlist_for_each_entry(sock_node, node, head, hlist_node) {
            inet_ntop(AF_INET, &sock_node->key.ip[0], addr, sizeof(addr));
            printf("sock_node:\n");
            printf("    queue_id : %d\n", sock_node->queue_id);
            printf("    cur_conn : %d\n", sock_node->cur_conn);
            printf("    ip       : %s\n", addr);
            printf("    port     : %d\n", ntohs(sock_node->key.port));
            printf("    proto    : %d\n", sock_node->key.proto);
        }
    }
    rte_rwlock_read_unlock(&sock_htable->rwlock);
}
